#ifndef SD_GMM_HPP
#define SD_GMM_HPP

#include <vector>
#include "real.hpp"

namespace MT {

// a point is a row-major vector

// -----------------------------------------------------------------------------------------
// distribution -- a distribution on a D-dimensional Euclidean space
// -----------------------------------------------------------------------------------------

// a distribution on a Euclidean space
// TODO: UnitTest
class distribution
{
protected: // internal data
    bool valid; // whether the distribution is valid
    int D; // dimensionality
    real weight; // scaling factor

public:
    distribution() : valid(false) {}

    // copying
    distribution(distribution const &inst) { *this = inst; }
    distribution &operator=(distribution const &inst)
    {
        valid = inst.valid;
        D = inst.D;
        weight = inst.weight;
        return *this;
    }

    bool is_valid() const { return valid; }
    void set_valid_state(bool valid) { this->valid = valid; }

    int get_dim() const { return D; }
    void set_dim(int D) { if(this->D != D) { this->D = D; valid = false; } }

    real get_weight() const { return weight; }
    void set_weight(real w)
    {
        CV_Assert(weight >= 0);
        this->weight = weight;
    }

    // compute the probability at a given point
    real operator()(cv::Mat const &x) { return 0; }

    // multiplication by a positive factor
    distribution &operator *=(real v) { set_weight(weight*v); return *this; }
    distribution operator *(real v) { distribution g(*this); g.set_weight(g.weight*v); return g; }

public: // boost::serialization
#ifdef USE_BOOST_SERIALIZATION
    friend class boost::serialization::access;
    template<class Archive>
    void serialize(Archive & ar, const unsigned int version) const
    {
        ar & BOOST_SERIALIZATION_NVP(valid) & BOOST_SERIALIZATION_NVP(D)
            & BOOST_SERIALIZATION_NVP(weight);
    }

#endif

};


// -----------------------------------------------------------------------------------------
// gaussian -- a multivariate gaussian distribution
// -----------------------------------------------------------------------------------------

// a multivariate gaussian distribution
// TODO: UnitTest
class gaussian : public distribution
{
protected: // internal data
    cv::Mat mean, cov, invCov;
    real detCov; // determinant of covariance
    real Z; // normalisation factor

public:
    gaussian() : distribution() {}

    gaussian(cv::Mat const &mean, cv::Mat const &covariance, real weight=1.0,
        cv::Mat const &_invCov=cv::Mat(), real _detCov=0)
    {
        set(mean, covariance, weight, _invCov, detCov);
    }

    // copying
    gaussian(gaussian const &inst) { *this = inst; }
    gaussian &operator=(gaussian const &inst)
    {
        distribution::operator=(inst);
        inst.mean.copyTo(mean);
        inst.cov.copyTo(cov);
        invCov = inst.invCov;
        detCov = inst.detCov;
        Z = inst.Z;
        return *this;
    }

    cv::Mat const &get_mean() const { return mean; }
    void set_mean(cv::Mat const &mean)
    {
        CV_Assert(mean.rows == 1 && mean.cols == D);
        mean.copyTo(this->mean);
    }

    real get_normalisation_factor() const { return Z; }
    cv::Mat const &get_covariance() const { return cov; }
    cv::Mat const &get_invCov() const { return invCov; }
    real get_detCov() const { return detCov; }
    void set_covariance(cv::Mat const &covariance, cv::Mat const &_invCov=cv::Mat(), real _detCov=0)
    {
        CV_Assert(covariance.rows == D && covariance.cols == D);
        covariance.copyTo(cov);

        if(_invCov.empty()) detCov = cv::invert(covariance, invCov);
        else
        {
            _invCov.copyTo(invCov);
            detCov = _detCov > 0? _detCov: cv::determinant(covariance);
        }
        CV_Assert(detCov > 0);

        Z = 1.0 / (pow(M_PI*2, 0.5*D)*sqrt(detCov));
    }

    real get_Z() const { return Z; }

    void set(cv::Mat const &mean, cv::Mat const &covariance, real weight=1.0,
        cv::Mat const &_invCov=cv::Mat(), real _detCov=0)
    {
        set_dim(mean.cols);
        set_weight(weight);
        set_mean(mean);
        set_covariance(covariance, _invCov, _detCov);
        set_valid_state(true);
    }


    // compute the entropy of the distribution
    real entropy() { CV_Assert(valid); return 0.5*D - log(Z); }

    // compute the probability at a given point
    real operator()(cv::Mat const &x)
    {
        CV_Assert(valid);
        real t;
        cv::Mat y, z;
        cv::subtract(x, mean, y);
        cv::gemm(y, invCov, 1, cv::Mat(), 0, z, cv::GEMM_2_T);
        t = y.dot(z);
        if(t > 100) return 0; // too small
        return Z*exp(-0.5*t);
    }

    // linear transform: y = Ax+b -- (y, x and b are row-major vectors)
    gaussian linear_transform(cv::Mat const &A, cv::Mat const &b)
    {
        CV_Assert(valid && b.rows==1);
        int D2 = b.cols;
        CV_Assert(A.rows==D2 && A.cols==D);
        cv::Mat mean_y, cov_y;

        // covariance
        cv::gemm(cov, A, 1, cv::Mat(), 0, mean_y, cv::GEMM_2_T); // Cov A'
        cv::gemm(A, mean_y, 1, cv::Mat(), 0, cov_y); // A Cov A'

        // mean
        cv::gemm(mean, A, 1, b, 1, mean_y, cv::GEMM_2_T);

        return gaussian(mean_y, cov_y, weight);
    }

    // linear project: y = w' x -- (y, x and w are row-major vectors)
    gaussian linear_project(cv::Mat const &w)
    {
        CV_Assert(valid && w.rows==1 && w.cols==D);
        cv::Mat mean_y, cov_y;

        // covariance
        cv::gemm(w, cov, 1, cv::Mat(), 0, mean_y); // w Cov
        cv::gemm(mean_y, w, 1, cv::Mat(), 0, cov_y, cv::GEMM_2_T); // w Cov w'

        // mean
        cv::gemm(w, mean, 1, cv::Mat(), 0, mean_y, cv::GEMM_2_T);

        return gaussian(mean_y, cov_y, weight);
    }

public: // boost::serialization
#ifdef USE_BOOST_SERIALIZATION
    friend class boost::serialization::access;
    template<class Archive>
    void serialize(Archive & ar, const unsigned int version) const
    {
        ar & boost::serialization::base_object<distribution>(*this);

        cv::Mat old_cov = cov;
        ar & BOOST_SERIALIZATION_NVP(mean) & BOOST_SERIALIZATION_NVP(cov);
        if(cov.data != old_cov.data) set_covariance(cov);
    }

#endif

};


// -----------------------------------------------------------------------------------------
// gmm -- a gaussian mixture model
// -----------------------------------------------------------------------------------------

// a gaussian mixture model
// TODO: UnitTest
class gmm : public distribution
{
public: // internal data
    std::vector<gaussian> g;

public:
    gmm() : distribution() {}
    gmm(std::vector<gaussian> const &g) { set(g); }
    gmm(std::vector<gaussian const *> const &g) { set(g); }

    // copying
    gmm(gmm const &inst) { *this = inst; }
    gmm &operator=(gmm const &inst)
    {
        distribution::operator=(inst);
        g = inst.g;
        return *this;
    }

    void set(std::vector<gaussian> const &_g)
    {
        int n = _g.size();
        std::vector<gaussian const *> _g2(n);
        while(--n >= 0) _g2[n] = &_g[n];
        set(_g2);
    }

    void set(std::vector<gaussian const *> const &_g)
    {
        real sw = 0;
        int i, N = _g.size();
        for(i = 0; i < N; ++i)
        {
            CV_Assert(_g[i]->is_valid());
            sw += _g[i]->get_weight();
        }
        g.resize(N);
        for(i = 0; i < N; ++i)
        {
            g[i] = *_g[i];
            g[i].set_weight(_g[i]->get_weight() / sw);
        }
        set_weight(sw);
        set_valid_state(true);
    }

    // compute the probability at a given point
    real operator()(cv::Mat const &x)
    {
        CV_Assert(valid);
        real v = 0;
        int n = g.size(); while(--n >= 0) v += g[n].get_weight() * g[n](x);
        return v / weight;
    }

    // linear transform: y = Ax+b -- (y, x and b are row-major vectors)
    gmm linear_transform(cv::Mat const &A, cv::Mat const &b)
    {
        int n = g.size();
        std::vector<gaussian> v(n);
        while(--n >= 0) v[n] = g[n].linear_transform(A, b);
        return gmm(v);
    }

    // linear project: y = w' x -- (y, x and w are row-major vectors)
    gmm linear_project(cv::Mat const &w)
    {
        int n = g.size();
        std::vector<gaussian> v(n);
        while(--n >= 0) v[n] = g[n].linear_project(w);
        return gmm(v);
    }

public: // boost::serialization
#ifdef USE_BOOST_SERIALIZATION
    friend class boost::serialization::access;
    template<class Archive>
    void serialize(Archive & ar, const unsigned int version) const
    {
        ar & boost::serialization::base_object<distribution>(*this);
        ar & BOOST_SERIALIZATION_NVP(g);
    }

#endif

};


// -----------------------------------------------------------------------------------------
// gaussian/gmm operations
// -----------------------------------------------------------------------------------------


// sum of a gaussian/gmm and a gaussian/gmm
gmm operator +(gaussian const &x, gaussian const &y);
gmm operator +(gaussian const &x, gmm const &y);
gmm operator +(gmm const &x, gaussian const &y);
gmm operator +(gmm const &x, gmm const &y);

// product of a gaussian/gmm and a gaussian/gmm
gaussian operator *(gaussian const &x, gaussian const &y);
gmm operator *(gaussian const &x, gmm const &y);
gmm operator *(gmm const &x, gaussian const &y);
gmm operator *(gmm const &x, gmm const &y);

// convolution of a gaussian/gmm and a gaussian/gmm
gaussian convolve(gaussian const &x, gaussian const &y);
gmm convolve(gaussian const &x, gmm const &y);
gmm convolve(gmm const &x, gaussian const &y);
gmm convolve(gmm const &x, gmm const &y);


// reduce the number of gaussians by grouping clusters of gaussians
//
// Gaussians are grouped into clusters by running weighted K-Means clustering
// on top of their centers. Then, each cluster of gaussians is approximated
// by a single gaussian using maximum likelihood
gmm reduce(gmm const &inst, int K);

// reduce the gmm down to a single gaussian by using an unbiased estimation
gaussian reduce(gmm const &inst);

} // namespace MT

#endif
